package net.aiion.weave.impl.messaging.transformation;

import java.util.ArrayList;
import java.util.Collection;

import net.aiion.weave.asm.Label;
import net.aiion.weave.asm.Opcodes;
import net.aiion.weave.asm.Type;
import net.aiion.weave.spi.bytecode.Bytecode;
import net.aiion.weave.spi.bytecode.BytecodeContext;
import net.aiion.weave.spi.bytecode.BytecodeException;
import net.aiion.weave.spi.bytecode.BytecodeGenerator;
import net.aiion.weave.spi.bytecode.events.AnnotationEndBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.AnnotationStartBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.AnnotationValueBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.BytecodeEvent;
import net.aiion.weave.spi.bytecode.events.ClassEndBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.ClassStartBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.MethodEndBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.MethodStartBytecodeEvent;
import net.aiion.weave.spi.bytecode.util.BytecodeTransformation;
import net.aiion.weave.spi.messaging.MessageReceiver;
import net.aiion.weave.spi.messaging.MessageReceiverContext;
import net.aiion.weave.spi.messaging.annotations.Receives;

// TODO Add inheritance support (also respect annotations of implemented interfaces).
public class MessageReceiverTransformation extends BytecodeTransformation {
    private static final String MESSAGE_RECEIVER_CONTEXT_INTERNAL_NAME = Type
        .getInternalName(MessageReceiverContext.class);

    private Collection<BytecodeEvent> receiveMessageBytecodeEvents;
    private BytecodeGenerator receiveMessageBytecodeGenerator;

    private String className;

    private String methodName;
    private String methodDescriptor;
    private String queueName;

    private boolean inReceivesAnnotation;

    public MessageReceiverTransformation() {
    }

    @Override
    public boolean testRetain(final BytecodeContext aBytecodeContext, final BytecodeEvent aBytecodeEvent) {
        return aBytecodeEvent.getType() == BytecodeEvent.CLASS_START
            && !aBytecodeEvent.asClassStart().containsInterface(MessageReceiver.class);
    }

    @Override
    public boolean testProcess(final BytecodeContext aBytecodeContext, final BytecodeEvent aBytecodeEvent) {
        return aBytecodeEvent.getType() == BytecodeEvent.ANNOTATION_START
            && aBytecodeEvent.asAnnotationStart().equalsType(Receives.class);
    }

    @Override
    public boolean testSkip(final BytecodeContext aBytecodeContext, final BytecodeEvent aBytecodeEvent) {
        return aBytecodeEvent.getType() == BytecodeEvent.CLASS_START
            && (aBytecodeEvent.asClassStart().isInterface() || aBytecodeEvent.asClassStart().containsInterface(
                MessageReceiver.class));
    }

    @Override
    public void handleBytecodeEvent(final BytecodeContext aBytecodeContext, final BytecodeEvent aBytecodeEvent) {
        switch (aBytecodeEvent.getType()) {
        case BytecodeEvent.CLASS_START:
            handleClassStartBytecodeEvent(aBytecodeContext, aBytecodeEvent.asClassStart());
            break;
        case BytecodeEvent.METHOD_START:
            handleMethodStartBytecodeEvent(aBytecodeContext, aBytecodeEvent.asMethodStart());
            break;
        case BytecodeEvent.ANNOTATION_START:
            handleAnnotationStartBytecodeEvent(aBytecodeContext, aBytecodeEvent.asAnnotationStart());
            break;
        case BytecodeEvent.ANNOTATION_VALUE:
            handleAnnotationValueBytecodeEvent(aBytecodeContext, aBytecodeEvent.asAnnotationValue());
            break;
        case BytecodeEvent.ANNOTATION_END:
            handleAnnotationEndBytecodeEvent(aBytecodeContext, aBytecodeEvent.asAnnotationEnd());
            break;
        case BytecodeEvent.METHOD_END:
            handleMethodEndBytecodeEvent(aBytecodeContext, aBytecodeEvent.asMethodEnd());
            break;
        case BytecodeEvent.CLASS_END:
            handleClassEndBytecodeEvent(aBytecodeContext, aBytecodeEvent.asClassEnd());
            break;
        default:
            break;
        }

        aBytecodeContext.forwardBytecodeEvent(aBytecodeEvent);
    }

    private void handleClassStartBytecodeEvent(final BytecodeContext aBytecodeContext,
        final ClassStartBytecodeEvent aBytecodeEvent) {
        className = aBytecodeEvent.getName();
        aBytecodeEvent.addInterface(MessageReceiver.class);

        prepareReceiveMessage();
    }

    private void handleMethodStartBytecodeEvent(final BytecodeContext aBytecodeContext,
        final MethodStartBytecodeEvent aBytecodeEvent) {
        methodName = aBytecodeEvent.getName();
        methodDescriptor = aBytecodeEvent.getDescriptor();

        if (Type.getArgumentTypes(methodDescriptor).length != 1) {
            throw new BytecodeException(String.format(
                "A @Receives annotated method must accept exactly one argument of type %s or any message type",
                MessageReceiverContext.class));
        }
    }

    private void handleAnnotationStartBytecodeEvent(final BytecodeContext aBytecodeContext,
        final AnnotationStartBytecodeEvent aBytecodeEvent) {
        if (methodName != null && aBytecodeEvent.equalsType(Receives.class)) {
            inReceivesAnnotation = true;
        }
    }

    private void handleAnnotationValueBytecodeEvent(final BytecodeContext aBytecodeContext,
        final AnnotationValueBytecodeEvent aBytecodeEvent) {
        if (inReceivesAnnotation && "value".equals(aBytecodeEvent.getName())) {
            queueName = (String) aBytecodeEvent.getValue();
        }
    }

    private void handleAnnotationEndBytecodeEvent(final BytecodeContext aBytecodeContext,
        final AnnotationEndBytecodeEvent aBytecodeEvent) {
        inReceivesAnnotation = false;
    }

    private void handleMethodEndBytecodeEvent(final BytecodeContext aBytecodeContext,
        final MethodEndBytecodeEvent aBytecodeEvent) {
        if (queueName != null) {
            extendReceiveMessage();
        }

        methodName = null;
        queueName = null;
    }

    private void handleClassEndBytecodeEvent(final BytecodeContext aBytecodeContext,
        final ClassEndBytecodeEvent aBytecodeEvent) {
        finishReceiveMessage();

        for (BytecodeEvent bytecodeEvent : receiveMessageBytecodeEvents) {
            aBytecodeContext.forwardBytecodeEvent(bytecodeEvent);
        }
    }

    private void prepareReceiveMessage() {
        receiveMessageBytecodeEvents = new ArrayList<>();
        receiveMessageBytecodeGenerator = Bytecode.createGenerator(receiveMessageBytecodeEvents);

        receiveMessageBytecodeGenerator.startMethod(Opcodes.ACC_PUBLIC, "receiveMessage", void.class,
            MessageReceiverContext.class);
        receiveMessageBytecodeGenerator.startCode();
    }

    private void extendReceiveMessage() {
        final Label nextReceiverLabel = new Label();

        final Type argumentType = Type.getArgumentTypes(methodDescriptor)[0];

        receiveMessageBytecodeGenerator.pushConstant(queueName);
        receiveMessageBytecodeGenerator.loadReference(1);
        receiveMessageBytecodeGenerator.invokeInterface(MessageReceiverContext.class, "getQueueName", String.class);
        receiveMessageBytecodeGenerator.invokeVirtual(String.class, "equals", boolean.class, Object.class);
        receiveMessageBytecodeGenerator.jumpEquals(nextReceiverLabel);

        if (!argumentType.getInternalName().equals(MESSAGE_RECEIVER_CONTEXT_INTERNAL_NAME)) {
            receiveMessageBytecodeGenerator.pushConstant(argumentType);
            receiveMessageBytecodeGenerator.loadReference(1);
            receiveMessageBytecodeGenerator.invokeInterface(MessageReceiverContext.class, "getMessage", Object.class);
            receiveMessageBytecodeGenerator.invokeVirtual(Object.class, "getClass", Class.class);
            receiveMessageBytecodeGenerator.invokeVirtual(Class.class, "isAssignableFrom", boolean.class, Class.class);
            receiveMessageBytecodeGenerator.jumpEquals(nextReceiverLabel);
        }

        if (!methodDescriptor.endsWith(")V")) {
            receiveMessageBytecodeGenerator.loadReference(1);
        }

        receiveMessageBytecodeGenerator.loadReference(0);
        receiveMessageBytecodeGenerator.loadReference(1);

        if (!argumentType.getInternalName().equals(MESSAGE_RECEIVER_CONTEXT_INTERNAL_NAME)) {
            receiveMessageBytecodeGenerator.invokeInterface(MessageReceiverContext.class, "getMessage", Object.class);
            receiveMessageBytecodeGenerator.castChecked(argumentType);
        }
        receiveMessageBytecodeGenerator.invokeVirtual(className, methodName, methodDescriptor);

        if (!methodDescriptor.endsWith(")V")) {
            receiveMessageBytecodeGenerator.invokeInterface(MessageReceiverContext.class, "setResult", void.class,
                Object.class);
        }

        receiveMessageBytecodeGenerator.markLabel(nextReceiverLabel);
        receiveMessageBytecodeGenerator.markFrameSame();
    }

    private void finishReceiveMessage() {
        receiveMessageBytecodeGenerator.returnVoid();
        receiveMessageBytecodeGenerator.endCode(3, 2);
        receiveMessageBytecodeGenerator.endMethod();
    }
}
